# 10/100 子串-和为K的子数组
# leetcode第560题: https://leetcode.cn/problems/subarray-sum-equals-k/description/?envType=study-plan-v2&envId=top-100-liked
# Date: 2024/11/19
import json

from leetcode import test


def subarraySum_bf(nums: list[int], k: int) -> int:
    """暴力法, 超时"""
    n = len(nums)
    if n < 2:
        return 1 if nums[0] == k else 0

    res = 0
    for i in range(n):
        tem_sum = 0
        for j in range(i, n):
            tem_sum += nums[j]
            if tem_sum == k:
                res += 1

    return res


def subarraySum_prefixsum(nums: list[int], k: int) -> int:
    """使用一维动态规划, 仍然超时.
    """
    n = len(nums)
    if n < 2:
        return 1 if nums[0] == k else 0

    res = 0
    dp = [float('inf')] * n  # dp[i] 表示前i个数的和, 实际上是存储前缀和的数组
    dp[0] = nums[0]
    if nums[0] == k:
        res += 1
    for i in range(1, n):
        dp[i] = dp[i - 1] + nums[i]
        t = dp[i] - k
        if dp[i] == k:
            res += 1
        res += [c == t for c in dp[:i]].count(True)
    return res


def subarraySum_prefixsum_mapping(nums: list[int], k: int) -> int:
    """根据上面的思路, 我们进一步提出优化, 使用一个mapping来代替dp, 来优化下面这段代码
    res += [c == t for c in dp[:i]].count(True)
    """
    n = len(nums)
    if n < 2:
        return 1 if nums[0] == k else 0

    res = 0
    mapping = {}
    pre_sum = 0
    for i in range(n):
        pre_sum += nums[i]
        t = pre_sum - k
        if t == 0:
            res += 1

        # if mapping.get(t):
        #     res += mapping[t]
        res += mapping.get(t, 0)  # 上面代码的简化的代码

        # if mapping.get(pre_sum):
        #     mapping[pre_sum] += 1
        # else:
        #     mapping[pre_sum] = 1
        mapping[pre_sum] = mapping.get(pre_sum, 0) + 1  # 上面代码的简化

    return res


def subarraySum_prefixsum_mapping_opt(nums: list[int], k: int) -> int:
    """进一步优化代码
    """
    res = 0
    mapping = {0: 1}
    pre_sum = 0
    for n in nums:
        pre_sum += n
        t = pre_sum - k
        res += mapping.get(t, 0)
        mapping[pre_sum] = mapping.get(pre_sum, 0) + 1
    return res


if __name__ == '__main__':
    with open("../data/d_substr1_lc560.data", "r") as file:
        lines = file.read().strip()  # 去除首位的空白
        jo = json.loads(lines)

    inp = [{"nums": [1, 1, 1], "k": 2},
           {"nums": [1, 2, 3], "k": 3},
           {"nums": [1, 2, 1, 2, 1], "k": 3},
           {"nums": [-1, 2, 2, 1, 2, 3], "k": 3},
           {"nums": [-1, -1, 1], "k": 1},
           {"nums": jo["nums"], "k": jo["k"]}, ]
    out = [2, 2, 4, 4, 1, 54]
    test.test_function(subarraySum_prefixsum, inp, out)
    test.test_function(subarraySum_prefixsum_mapping, inp, out)
    test.test_function(subarraySum_prefixsum_mapping_opt, inp, out)
